import os
import json
import pandas as pd
import step4_check_batch, step5_parse_jsonl, step6_check_symbols, step7_merge_results
from utils import save_dataframe


def load_job_info(job_info_path):
    """加载任务信息文件"""
    if not os.path.exists(job_info_path):
        print(f"错误: 任务信息文件不存在: {job_info_path}")
        return None

    with open(job_info_path, 'r', encoding='utf-8') as f:
        job_info = json.load(f)

    print(f"加载任务信息: {job_info_path}")
    print(f"任务ID: {job_info['batch_id']}")
    print(f"输入文件: {job_info['input_excel']}")

    return job_info


def check_and_process(job_info):
    """检查任务状态并处理结果"""
    print("\n步骤4: 检查任务状态...")
    result_file = step4_check_batch.check_batch_status(job_info["batch_id"])
    print(f"结果文件: {result_file}")

    print("\n步骤5: 解析JSONL结果...")
    text_results = step5_parse_jsonl.parse_jsonl_results(result_file)
    save_dataframe(text_results, "文本校对结果.xlsx")

    print("\n步骤6: 校对符号内容...")
    # 加载符号内容
    symbol_df = pd.read_excel(job_info["symbol_file"])
    symbol_results = step6_check_symbols.check_symbol_content(symbol_df)
    save_dataframe(symbol_results, "符号校对结果.xlsx")

    print("\n步骤7: 合并结果...")
    # 加载完整文档
    full_df = pd.read_excel(job_info["full_file"])
    final_df = step7_merge_results.merge_results(full_df, text_results, symbol_results)
    final_path = save_dataframe(final_df, "最终校对结果.xlsx")

    print("\n处理完成!")
    print(f"最终结果已保存至: {final_path}")
    print(f"总行数: {len(final_df)}")
    print(f"通过: {len(final_df[final_df['校对状态'] == '通过'])}")
    print(f"未通过: {len(final_df[final_df['校对状态'] == '未通过'])}")

    # 更新任务状态
    job_info["status"] = "completed"
    job_info["result_file"] = os.path.abspath(final_path)
    with open(job_info_path, 'w', encoding='utf-8') as f:
        json.dump(job_info, f, ensure_ascii=False, indent=2)

    return final_path


if __name__ == "__main__":
    # 获取所有任务信息文件
    job_files = [f for f in os.listdir() if f.startswith('job_') and f.endswith('.json')]

    if not job_files:
        print("错误: 未找到任务信息文件")
        print("请先运行 create_batch_job.py 创建任务")
        exit(1)

    # 显示可用任务
    print("可用任务列表:")
    for i, job_file in enumerate(job_files, 1):
        print(f"{i}. {job_file}")

    # 选择任务
    try:
        choice = int(input("\n请选择要处理的任务编号: ")) - 1
        job_info_path = job_files[choice]
    except (ValueError, IndexError):
        print("无效的选择")
        exit(1)

    # 处理任务
    job_info = load_job_info(job_info_path)
    if job_info:
        if job_info.get("status") == "completed":
            print("\n警告: 此任务已完成处理")
            print(f"结果文件: {job_info.get('result_file', '未知')}")
            choice = input("是否重新处理? (y/n): ").lower()
            if choice != 'y':
                exit(0)

        check_and_process(job_info)